package com.zyx.flinkdemo.stream.statebackend;

import org.apache.flink.configuration.ConfigOption;
import org.apache.flink.configuration.ConfigOptions;
import org.apache.flink.configuration.ReadableConfig;
import org.apache.flink.contrib.streaming.state.ConfigurableRocksDBOptionsFactory;
import org.apache.flink.contrib.streaming.state.RocksDBOptionsFactory;
import org.rocksdb.BlockBasedTableConfig;
import org.rocksdb.ColumnFamilyOptions;
import org.rocksdb.DBOptions;

import java.util.Collection;

/**
 * @author zyx
 * @since 2021/8/14 00:29
 * desc: 配置rocksdb的参数类
 */
public class MyOptionsFactory implements ConfigurableRocksDBOptionsFactory {
    /**
     * 默认缓存大小为 256MB
     */
    private static final long DEFAULT_SIZE = 256 * 1024 * 1024;
    private long blockCacheSize = DEFAULT_SIZE;

    @Override
    public DBOptions createDBOptions(DBOptions currentOptions, Collection<AutoCloseable> handlesToClose) {
        return currentOptions.setIncreaseParallelism(4)
                .setUseFsync(false);
    }

    @Override
    public ColumnFamilyOptions createColumnOptions(
            ColumnFamilyOptions currentOptions, Collection<AutoCloseable> handlesToClose) {
        return currentOptions.setTableFormatConfig(
                new BlockBasedTableConfig()
                        .setBlockCacheSize(blockCacheSize)
                        // 128 KB
                        .setBlockSize(128 * 1024));
    }

    @Override
    public RocksDBOptionsFactory configure(ReadableConfig configuration) {
        ConfigOption<Long> blockCacheSizeConf = ConfigOptions
                .key("my.custom.rocksdb.block.cache.size").longType().defaultValue(blockCacheSize);
        this.blockCacheSize = configuration.get(blockCacheSizeConf);
        return this;
    }
}